from PIL import Image
from torchvision import transforms
from torch.utils.data import Dataset
import jsonlines

from torch import nn


class TextCartoonDataset(Dataset):
    def __init__(self,jsonl_file,size = 256 ,tokenizer=None):
        super(TextCartoonDataset,self).__init__()
        assert tokenizer is not None
        self.tokenizer = tokenizer
        self.size = size
        self.data = self.read_jsonl(jsonl_file)
        
        self.T = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])
        ])
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        source_image_path = self.data[index]['source_image']
        caption = self.data[index]['text']
        
        source_image = Image.open(source_image_path).convert('RGB').resize([self.size,self.size])
        input_ids = self.tokenizer(caption)
        
        return {
            'pixel_values':self.T(source_image),
            'input_ids':input_ids
        }
    
    def read_jsonl(self,file):
        data = []
        for line in jsonlines.open(file,'r'):
            data.append(line)
        return data